#!/usr/bin/env python

# Take data from the CurrentCost unit and update it directly to an rrd
# database and to a raw datafile.
# pir 2010/03/18

# Global variables
BASE = '/usr/local/power'    # A directory to store everything underneath.
PORT = '/dev/cu.usbserial'   # The port connected to the currentcost.
BAUD = 57600                 # 9600 for a classic, 57600 for a cc128.
RRD = '%s/rrd/cc-%%d.rrd' % BASE # rrdtool file, needs to be created first.
LOG = '%s/var/output.txt' % BASE # Plain text output file.
# The number of seconds to keep running without a successful read.
DIE_TIME = 30
# Number of fully invalid responses to ignore before dying.
DIE_COUNT = 10
# Delay, in seconds or fractions of a second, per loop.
LOOP_DELAY = 0.1
#
DEBUG = False
#DEBUG = True

import sys
# Where to find pycurrentcost.py
sys.path.append('%s/lib' % BASE)
import pycurrentcost

import logging
import logging.handlers
import os
import rrdtool
import signal
import syslog
import time


class TimeoutException(Exception):
  pass 

def timeout_handler(signum, frame):
  raise TimeoutException()


def RrdUpdate(rrdfile, updatestr):
  """Take an rrd file and a string to update to it."""
  if rrdfile:
    try:
      rrdtool.update(rrdfile, updatestr)
    except rrdtool.error as e:
      syslog.syslog(syslog.LOG_ERR, 'Error updating RRD %s: %s' % (rrdfile, e))

def WriteUpdate(sensor, watts, time, temp):
  """Write an update out to rrd/file/etc."""

  # If we got an update, push data out.
  if watts != None:

    # The meters really aren't that sensitive for small power amounts.
    if watts < 5:
      watts = 0

    updatestr = str(u'%d:%s:%s' % (time, watts, temp))
    logging.debug('Writing data from sensor %d: %s' % (sensor, updatestr))

    # Write it all to a flat text file.
    if LOG:
      logfile.write('%d %s\n' % (sensor, updatestr))
      logfile.flush()

    # Log to a specific rrd file
    if RRD:
      RrdUpdate(RRD % sensor, updatestr)

def InitSysLog(debug=False,
               facility=logging.handlers.SysLogHandler.LOG_DAEMON,
               format='%(filename)s[%(process)d]: %(levelname)s %(message)s'):
  """Set up the logging module to log to syslog and stderr."""

  platform = sys.platform
  if platform == "darwin":
    # Apple made 10.5 more secure by disabling network syslog:
    address = "/var/run/syslog"
  elif platform.startswith('linux') or platform.startswith('freebsd'):
    address = "/dev/log"
  else:
    address = ('localhost', 514)
  syslog = logging.handlers.SysLogHandler(address, facility=facility)

  logger = logging.getLogger()
  if debug:
    logger.setLevel(logging.DEBUG)
  else:
    logger.setLevel(logging.INFO)

  formatter = logging.Formatter(format)
  syslog.setFormatter(formatter)
  logger.addHandler(syslog)

  logstderr = logging.StreamHandler(sys.stderr)
  logstderr.setFormatter(formatter)
  logger.addHandler(logstderr)

if __name__ == '__main__':

  # Init syslog.
  InitSysLog(debug=DEBUG)

  # Init the CurrentCost unit.
  cc = pycurrentcost.CurrentCostReader(port=PORT, baudrate=BAUD)

  # Make sure our output is readable if it didn't exist already.
  os.umask(0022)
  # Open the plain text logfile.
  if LOG:
    logfile = open(LOG, 'a')

  # Set an initial update time.
  updatetime = int(time.time())

  # Keep track of when we last saw an update for a given sensor
  last_update = {}

  # Number of bad responses.
  bad_count = 0

  # Loop and take readings.
  while True:

    watts = None

    # Timeout the command.
    old_handler = signal.signal(signal.SIGALRM, timeout_handler) 
    signal.alarm(12) # trigger alarm in 12 seconds

    try:
      reading = cc.get_reading()

    except pycurrentcost.CurrentCostException as e:
      logging.error('Error from get_reading: %s' % e)
    except pycurrentcost.CurrentCostXMLParseException as e:
      logging.error('Parse error from get_reading: %s' % e)
    except pycurrentcost.CurrentCostReadException as e:
      logging.debug('Read error from get_reading: %s' % e)
    except TimeoutException:
      logging.error('Timeout in get_reading')

    else:
      if reading:
        updatetime = int(time.time())

    finally:
      signal.signal(signal.SIGALRM, old_handler) 

    # Cancel the timeout.
    signal.alarm(0)

    # If we didn't get a defined reading out of the above.
    try:
      sensor_num = int(reading.sensor_num)
    except (NameError, TypeError) as e:
      logging.error('reading got an invalid response: %e' % e)
      bad_count += 1
      if bad_count > DIE_COUNT:
        logging.error('bad reading count exceeded %s, exiting' % DIE_COUNT)
        sys.exit(2)

    logging.debug('Sensor %d: channel data: %s' % (sensor_num,
                                                   reading.channels))

    # Example for fixed sensor/channel mapping.
    #if (1 in reading.channels):
    #  WriteUpdate(sensor_num, int(reading.channels[2]['watts']),
    #              updatetime, reading.temperature)
    # Map sensor 0 channel 2 -> sensor 4.
    #if (sensor_num == 0 and 2 in reading.channels):
    #  WriteUpdate(4, int(reading.channels[2]['watts']),
    #              updatetime, reading.temperature)

    # If we have have more than one channel in use then map it to
    # a sensor number greater than 9 which doesn't physically exist.
    # Sensor 0, channel 2 -> sensor 10
    # Sensor 1, channel 2 -> sensor 12
    # Sensor 1, channel 3 -> sensor 13
    # etc, up to sensor 29

    for channel in reading.channels:
      if (channel == "sensor_num"):
        continue
      channel = int(channel)
      if (channel != 1):
        sensor_num += sensor_num + 8 + channel

      sensor_index = '%s' % sensor_num
      if sensor_index not in last_update:
        # no recorded update for this sensor
        last_update[sensor_index] = 0

      if last_update[sensor_index] < updatetime:
        # We haven't updated this sensor this second
        last_update[sensor_index] = updatetime
        WriteUpdate(sensor_num, reading.channels[channel]['watts'],
                    updatetime, reading.temperature)

    # Check if we've had a recent update or not.
    if updatetime + DIE_TIME <= int(time.time()):
      logging.error('No updates for %s seconds, suiciding.' % DIE_TIME)
      sys.exit(1)

    # Pause enough time to not be spinning endlessly and taking up too
    # much CPU time but still get a reading per second from each sensor.
    time.sleep(LOOP_DELAY)
